import lightgbm as lgb
import pandas as pd
import json
import argparse
import shap
import numpy as np
def predict(modelpath,filepath,lgb_feature_importance_path,shap_feature_importance_path):
    gbm=lgb.Booster(model_file=modelpath)
    gbm.params['objective'] = 'multiclass'
    feature_name_list=gbm.feature_name()
    importance_list=gbm.feature_importance().tolist()
    feature_importancedict=dict(zip(feature_name_list,importance_list))

    feature_importancejson=json.dumps(feature_importancedict,ensure_ascii=False)
    # print(feature_importancejson)
    with open(lgb_feature_importance_path,'w',encoding='utf-8') as writejson:
        writejson.write(feature_importancejson)
    if(filepath!=''):
        data=pd.read_csv(filepath)
        explainer = shap.TreeExplainer(gbm)
        shap_values = explainer.shap_values(data[feature_name_list])  # 传入特征矩阵X，计算SHAP值
        shap_values_arr=np.array(shap_values)
        s=np.transpose(shap_values_arr,(2,0,1))
        shap_feature_impotance=np.mean(np.abs(s),-1)
        print(shap_feature_impotance.shape)
        feature_dict={}
        for i in range(len(feature_name_list)):
            feature_dict[feature_name_list[i]]=shap_feature_impotance[i].tolist()
        print(feature_dict)
        shap_feature_importancejson=json.dumps(feature_dict,ensure_ascii=False)
        # print(feature_importancejson)
        with open(shap_feature_importance_path,'w',encoding='utf-8') as writejson:
            writejson.write(shap_feature_importancejson)

if __name__=='__main__':
    # predict('./model_2023-07-10_22-14-26.016396.txt','./data/test_2000_x.csv','lgb_feature_importance.json','shap_feature_importance.json')
    parser = argparse.ArgumentParser()
    parser.add_argument("modelpath",default='',type=str)
    parser.add_argument("filepath",default='',type=str)
    parser.add_argument("lgb_feature_importance_path",default='',type=str)
    parser.add_argument("shap_feature_importance_path",default='',type=str)
    args = parser.parse_args()
    predict(args.modelpath,args.filepath,args.lgb_feature_importance_path,args.shap_feature_importance_path)
